import signal
import sys
import cv2
import rospy
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
import message_filters
import pathlib

IMGAE_TOPIC = "/camera/color/image_raw"
DEPTH_TOPIC = "/camera/depth/image_raw"
FRAME_HZ = 30
IMAGE_SIZE = (640, 480)


def signal_handler(sig, frame):
    print("接收到信号，程序即将退出...")
    sys.exit(0)


signal.signal(signal.SIGINT, signal_handler)  # 注册 Ctrl+C 信号处理函数


class RGBRecorder:
    def __init__(self, video_path, video_max_secs=10):
        self.cv_bridge = CvBridge()
        self.video_path = video_path
        self.video_max_secs = video_max_secs
        self.init_flag = False
        self.file_index = 0
        self.initialize_video_writers()
        rospy.Subscriber(IMGAE_TOPIC, Image, self.rgb_callback)

    def __del__(self):
        # Make sure to release the last set of video writers
        self.out_rgb.release()

    def initialize_video_writers(self):
        self.file_index += 1
        print(self.file_index)
        rgb_filename = str((self.video_path / f'./output_rgb_{self.file_index}.mp4'))
        self.fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.out_rgb = cv2.VideoWriter(rgb_filename, self.fourcc, FRAME_HZ, IMAGE_SIZE)

    def rgb_callback(self, rgb_data):
        if not self.init_flag:
            self.start_time = rgb_data.header.stamp
            self.init_flag = True
        try:
            cv_image = self.cv_bridge.imgmsg_to_cv2(rgb_data, desired_encoding="bgr8")
            self.out_rgb.write(cv_image)
        except Exception as e:
            rospy.logerr("RGB callback error: %s", e)

        # Check if 5 seconds have passed to switch to a new video file
        if (rgb_data.header.stamp - self.start_time).to_sec() > self.video_max_secs:
            self.start_time = rgb_data.header.stamp  # Reset start time
            self.out_rgb.release()  # Release current VideoWriter
            self.initialize_video_writers()  # Initialize new VideoWriters for the next segment

    def run(self):
        rospy.spin()


class RGBDRecorder:
    def __init__(self, video_path, video_max_secs=10):
        self.cv_bridge = CvBridge()
        self.init_flag = False
        self.file_index = 0
        self.video_path = video_path
        self.video_max_secs = video_max_secs
        self.initialize_video_writers()
        image_sub = message_filters.Subscriber(IMGAE_TOPIC, Image)
        depth_sub = message_filters.Subscriber(DEPTH_TOPIC, Image)
        time_syncronizer = message_filters.ApproximateTimeSynchronizer([image_sub, depth_sub], queue_size=10, slop=float(2.0 / float(FRAME_HZ)))
        time_syncronizer.registerCallback(self.sync_callback)

    def __del__(self):
        # Make sure to release the last set of video writers
        self.out_rgb.release()
        self.out_depth.release()

    def initialize_video_writers(self):
        self.file_index += 1
        print(self.file_index)
        rgb_filename = str((self.video_path / f'./output_rgb_{self.file_index}.mp4'))
        depth_filename = str((self.video_path / f'./output_depth_{self.file_index}.mp4'))

        self.fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.out_rgb = cv2.VideoWriter(rgb_filename, self.fourcc, FRAME_HZ, IMAGE_SIZE)
        self.out_depth = cv2.VideoWriter(depth_filename, self.fourcc, FRAME_HZ, IMAGE_SIZE)

    def sync_callback(self, rgb_data, depth_data):
        if not self.init_flag:
            self.start_time = rgb_data.header.stamp
            self.init_flag = True
        try:
            cv_rgb_image = self.cv_bridge.imgmsg_to_cv2(rgb_data, desired_encoding="bgr8")
            self.out_rgb.write(cv_rgb_image)
            cv_depth_image = self.cv_bridge.imgmsg_to_cv2(depth_data, desired_encoding="32FC1")
            depth_colored = cv2.applyColorMap(cv2.convertScaleAbs(cv_depth_image, alpha=0.03), cv2.COLORMAP_JET)
            self.out_depth.write(depth_colored)
        except Exception as e:
            rospy.logerr("Sync callback error: %s", e)

        # Check if 5 seconds have passed to switch to a new video file
        if (rgb_data.header.stamp - self.start_time).to_sec() > self.video_max_secs:
            self.start_time = rgb_data.header.stamp  # Reset start time
            self.out_rgb.release()  # Release current VideoWriter
            self.out_depth.release()  # Release current VideoWriter
            self.initialize_video_writers()  # Initialize new VideoWriters for the next segment

    def run(self):
        rospy.spin()


def main():
    rospy.init_node('rgbd_recorder')
    video_path = pathlib.Path(__file__).resolve().parent.parent / 'video'
    video_path.mkdir(parents=True, exist_ok=True)
    recorder = RGBRecorder(video_path)
    # recorder = RGBDRecorder(video_path)
    try:
        recorder.run()
    except rospy.ROSInterruptException:
        pass
    del recorder


if __name__ == '__main__':
    main()
